# 构建的数据类型
import torch
# 移入绘图库
import matplotlib.pyplot as plt
# 模型框架
from torch import nn

noise = 0.2
points_a = torch.normal(0.5, noise, (20, 2))
labels_a = torch.zeros((20, 1))
features_a = torch.concatenate([points_a, labels_a], dim=1)
points_b = torch.normal(1.5, noise, (20, 2))
labels_b = torch.ones((20, 1))
features_b = torch.concatenate([points_b, labels_b], dim=1)
features = torch.concatenate([features_a, features_b], dim=0)
indices = torch.randperm(features.shape[0])
features = features[indices]
# -----------------模型--------------------
model = nn.Sequential(
    nn.Linear(2, 1),  # 输入特征数量（x,y） 输出特征数量(label)
    nn.Sigmoid()  # sigmoid 函数介于0~1  >0.5就是0（A）分类 <0.5是1（B）分类
)
# -----------------损失函数--------------------
criterion = nn.MSELoss()
# -----------------优化器--------------------
# model.parameters() 模型中存在的所有的w,b(训练的参数)
# lr = leaning_rate 学习率（控制梯度（w,b）下降）
sgd = torch.optim.SGD(model.parameters(), lr=0.1)

# -----------------训练--------------------
epochs = 3000  # 训练次数
for epoch in range(epochs):
    sgd.zero_grad()  # 清空梯度缓存
    predict_labels = model(features[:, :2])  # 将x，y 作为输入放入模型中预测
    loss = criterion(predict_labels, features[:, 2:])  # 将预测的label与实际的label进行方差比较
    loss.backward()  # 反向传播
    sgd.step()  # 更新所有的w,b 的参数

    # 输出每一次的loss的值
    print(f"{epoch + 1} / {epochs} -- loss:{loss.item():.4f}")

# -----------------测试不一样的数据集-------------------
x = torch.linspace(0, 2, 20)
y = torch.linspace(0, 2, 20)
x, y = torch.meshgrid([x, y], indexing='ij')
x = x.reshape(400, 1)
y = y.reshape(400, 1)
test_features = torch.concatenate([x, y], dim=1)
# 预测这些测试的数据集
model.eval()  # 开启模型测试（不可以用于训练）
predict_test = model(test_features)
# 由于预测的结果是0~1 之间的数值， <0.5 是A簇类（0） >0.5 是B簇类（1）
predict_test[predict_test < 0.5] = 0
predict_test[predict_test >= 0.5] = 1


# -----------------绘制图线-------------------
def draw_points(input, output, ax):
    for i, (x, y) in enumerate(input):
        # print(x, y, output[i, 0])
        if output[i, 0].item() == 0:
            ax.plot([x.item()], [y.item()], 'ro')
        elif output[i, 0].item() == 1:
            ax.plot([x.item()], [y.item()], 'bo')


# 绘制一行两列的第一个图像
ax1 = plt.subplot(121)
# 绘制出原来簇点图
draw_points(features[:, :2], features[:, 2:], ax1)
# 绘制一行两列的第二个图像
ax2 = plt.subplot(122)
# 绘制测试集中的图（展示训练后的结果效果）
draw_points(test_features, predict_test, ax2)
plt.show()
